#include "hs_cuda_utils.cuh"
#include "hs_logger.h"

cudaError_t hs::cuda::showGPUInfo(int id)
{
    /* 总的设备数 */
    int num_device;
    {
        cudaError_t err = cudaGetDeviceCount(&num_device);
        if (err != cudaSuccess) {
            LOG_ERROR("cudaGetDeviceCount failed!\n");
            return err;
        }
    }

    if (id >= num_device) {
        LOG_WARN("`id` exceeds the number of devices.\n");
        id = 0;
    }
    
    cudaDeviceProp dev_prop;
    {
        cudaError_t err = cudaGetDeviceProperties(&dev_prop, id);
        if (err != cudaSuccess) {
            LOG_ERROR("cudaGetDeviceProperties failed!\n");
            return err;
        }
    }

    LOG_BUFFER_BEGIN
    LOG_BUFFER("GPU device [" $ id $ "]: " $ dev_prop.name $ std::endl);
    LOG_BUFFER("|---- compute Architecture Version: " $ dev_prop.major $ "." $ dev_prop.minor $ std::endl);
    LOG_BUFFER("|---- totalGlobalMemory: " $ dev_prop.totalGlobalMem / 1024.0 / 1024.0 $ " MB\n");
    LOG_BUFFER("|---- sharedMemPerBlock: " $ dev_prop.sharedMemPerBlock / 1024.0 $ " KB\n");
    LOG_BUFFER("|---- maxThreadsPerBlock: " $ dev_prop.maxThreadsPerBlock $ " \n");
    LOG_BUFFER("|---- regsPerBlock: " $ dev_prop.regsPerBlock $ "\n");
    LOG_BUFFER("|---- maxThreadsPerMultiProcessor: " $ dev_prop.maxThreadsPerMultiProcessor $ "\n");
    LOG_BUFFER("|____ multiProcessorCount: " $ dev_prop.multiProcessorCount $ " \n");
    LOG_BUFFER_END

    return cudaSuccess;
}

cudaError_t hs::cuda::showAllGPUInfo()
{
    /* 总的设备数 */
    int num_device;
    {
        cudaError_t err = cudaGetDeviceCount(&num_device);
        if (err != cudaSuccess) {
            LOG_ERROR("cudaGetDeviceCount failed!\n");
            return err;
        }
    }

    for (int i = 0; i < num_device; ++i) {
        cudaError_t err = showGPUInfo(i);
        if (err != cudaSuccess) {
            LOG_ERROR("showAllGPUInfo failed!\n");
            return err; 
        }
    }

    return cudaSuccess;
}

cudaError_t hs::cuda::GetNumBlocks(int n, int * num_blocks)
{
    /* 当前使用的设备号 */
    int dev_id;
    {
        cudaError err = cudaGetDevice(&dev_id);
        if (err != cudaSuccess) {
            return err;
        }
    }

    /* 获取多处理器个数 */
    int sm_count;
    {
        cudaError err = cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, dev_id);
        if (err != cudaSuccess) {
            return err;
        }
    }

    /* 获取每个多处理的最大线程数 */
    int tpm;
    {
        cudaError_t err = cudaDeviceGetAttribute(&tpm, cudaDevAttrMaxThreadsPerMultiProcessor, dev_id);
        if (err != cudaSuccess) {
            return err;
        }
    }

    *num_blocks = (n + KERNEL_BLOCK_SIZE - 1) / KERNEL_BLOCK_SIZE;
    int tmp1 = sm_count * tpm * KERNEL_NUM_WAVES / KERNEL_BLOCK_SIZE;
    *num_blocks = (tmp1 < (*num_blocks)) ? tmp1 : (*num_blocks);
    *num_blocks = (1 > (*num_blocks)) ? 1 : (*num_blocks);

    return cudaSuccess;
}
